from d435 import D435
import cv2
from ultralytics import YOLO
import numpy as np
import o3d_util as o3d
import cv_util as cvu

def yolomask_to_pointcloud(color_image, depth_image, yolo_mask):
    yolo_mask_float = yolo_mask.reshape(yolo_mask.shape[1], -1)
    print(yolo_mask.shape)
    size_image = (color_image.shape[1], color_image.shape[0])
    yolo_mask_resize= cv2.resize(yolo_mask_float, size_image, interpolation=cv2.INTER_LINEAR)
    yolo_mask_resize_uint8 = yolo_mask_resize.astype(np.uint8)
    yolo_mask_resize_uint8 = yolo_mask_resize_uint8*255
    color_image_masked = cv2.add(color_image, np.zeros(np.shape(color_image), dtype=np.uint8), mask=yolo_mask_resize_uint8)
    depth_image_masked = cv2.add(depth_image, np.zeros(np.shape(depth_image), dtype=np.uint16), mask=yolo_mask_resize_uint8)
    cvu.depth_filter(depth_image_masked)
    pcd = o3d.gen_point_cloud(color_image_masked, depth_image_masked)
    return pcd

if __name__ == '__main__':
    model_path = r'./yolov8n-seg.pt'
    model = YOLO(model_path)
    camera = D435()
    inr = camera.color_intrinsics
    while True:
        camera.update()
        color_image, depth_image = camera.get_images()
        cv2.imshow("detection", color_image)
        results = model(color_image)
        annotation_frame = results[0].plot()
        if results[0].masks is not None:
            mask = results[0].masks[0].data.cpu().numpy()
            pcd = yolomask_to_pointcloud(color_image, depth_image, mask)
            o3d.geometry_show([pcd])
            # out.write(annotation_frame)  # 写入帧
            cv2.imshow("A video", annotation_frame)
            # cv2.imshow("A video", color_image_masked)
            c = cv2.waitKey(30)
            if c == 27:
                break
